import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map.Entry;
import java.util.Set;

public class EmailReader {
	public HashMap<String, Integer> unigramlist;
	public HashMap<String, Integer> bigramlist;
	public ArrayList<String> authors;
	public ArrayList<String>[] unigramstringarr;
	public ArrayList<Integer>[] unigramvaluearr;
	public ArrayList<String>[] bigramstringarr;
	public ArrayList<Integer>[] bigramvaluearr;
	int[][] smoothedfreq;
	ArrayList<Integer> totalunifreq;

	public static void main(String[] args) {
		EmailReader reader = new EmailReader("Enrontrain.txt", "Encronvalidation.txt");
	}

	public EmailReader(String trainfile, String testfile) {
		unigramlist = new HashMap<String, Integer>();
		bigramlist = new HashMap<String, Integer>();
		totalunifreq = new ArrayList<Integer>();
		runAnalysis(trainfile);
		Extract();
		SmoothTuring();
		ArrayList<String> guesses = guessAuthors(testfile);
		try {
			FileWriter fwrite = new FileWriter("results.txt", false);
			BufferedWriter br = new BufferedWriter(fwrite);
			for (int i = 0; i < guesses.size(); i++) {
				br.write(guesses.get(i) + "\n");
			}
			for (int i = 0; i < 2024; i++) {
				br.write("derp\n");
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public void runAnalysis(String trainfile) {
		try {
			FileReader fread = new FileReader(trainfile);
			BufferedReader br = new BufferedReader(fread);
			String strLine;
			while ((strLine = br.readLine()) != null) {
				String[] array = strLine.split("\\s+");
				ArrayList<String> targetlist = new ArrayList<String>();
				boolean firstwordUnknown = false;
				for (int i = 0; i < array.length; i++) {
					targetlist.addAll(breakUpPunctuation(array[i]));
				}
				String author = targetlist.get(0);
				if (!authors.contains(author)) {
					authors.add(author);
				}
				if (!unigramlist.containsKey(author + " <UNK>")) {
					unigramlist.put(author + " <UNK>", 0);
				}
				if (targetlist.size() > 1) {
					String word = targetlist.get(1).toLowerCase();
					Integer freq = unigramlist.get(author + " " + word);
					if (freq == null) {
						unigramlist.put(author + " " + word, 0);
						word = "<UNK>";
						freq = unigramlist.get(author + " " + word);
						firstwordUnknown = true;
					}
					unigramlist.put(author + " " + word, freq + 1);
					totalunifreq.set(authors.indexOf(author), totalunifreq.get(authors.indexOf(author)) + 1);
				}
				for (int i = 2; i < targetlist.size(); i++) {
					String firstword = firstwordUnknown ? "<UNK>" : targetlist.get(i - 1).toLowerCase();
					String secondword = targetlist.get(i).toLowerCase();
					Integer freqUni = unigramlist.get(author + " " + secondword);
					if (freqUni == null) {
						unigramlist.put(author + " " + secondword, 0);
						secondword = "<UNK>";
						freqUni = unigramlist.get(author + " " + secondword);
						firstwordUnknown = true;
					} else {
						firstwordUnknown = false;
					}
					Integer freqBi = bigramlist.get(author + " " + firstword + " " + secondword);
					unigramlist.put(author + " " + secondword, freqUni + 1);
					bigramlist.put(author + " " + firstword + " " + secondword, freqBi != null ? freqBi + 1 : 1);
				}
			}
			// Close the input stream
			br.close();
		} catch (Exception e) {// Catch exception if any
			e.printStackTrace();
		}
	}

	private void Extract() {
		ArrayList<String> unigramstring = new ArrayList<String>();
		ArrayList<Integer> unigramvalue = new ArrayList<Integer>();
		ArrayList<String> bigramstring = new ArrayList<String>();
		ArrayList<Integer> bigramvalue = new ArrayList<Integer>();
		smoothedfreq = new int[authors.size()][5];
		// Get the data in usable form
		Set<Entry<String, Integer>> unigramset = unigramlist.entrySet();
		Set<Entry<String, Integer>> bigramset = bigramlist.entrySet();
		Iterator<Entry<String, Integer>> unigramit = unigramset.iterator();
		Iterator<Entry<String, Integer>> bigramit = bigramset.iterator();
		while (unigramit.hasNext()) {
			Entry<String, Integer> unigramentry = unigramit.next();
			if (unigramentry.getValue() != 0) {
				unigramstring.add(unigramentry.getKey());
				unigramvalue.add(unigramentry.getValue());
			}
		}
		while (bigramit.hasNext()) {
			Entry<String, Integer> bigramentry = bigramit.next();
			if (bigramentry.getValue() != 0) {
				bigramstring.add(bigramentry.getKey());
				bigramvalue.add(bigramentry.getValue());
			}
		}
		// We have extracted the data into arrays. now sort by authors
		for (int i = 0; i < unigramstring.size(); i++) {
			String[] s = unigramstring.get(i).split(" ");
			int index = authors.indexOf(s[0]);
			if (unigramstringarr[index] == null) {
				unigramstringarr[index] = new ArrayList<String>();
				unigramvaluearr[index] = new ArrayList<Integer>();
			}
			unigramstringarr[index].add(s[1]);
			unigramvaluearr[index].add(unigramvalue.get(i));
		}
		for (int i = 0; i < bigramstring.size(); i++) {
			String[] s = bigramstring.get(i).split(" ");
			int index = authors.indexOf(s[0]);
			if (bigramstringarr[index] == null) {
				bigramstringarr[index] = new ArrayList<String>();
				bigramvaluearr[index] = new ArrayList<Integer>();
			}
			bigramstringarr[index].add(s[1] + " " + s[2]);
			bigramvaluearr[index].add(bigramvalue.get(i));
		}
		// Data is now sorted by authors
	}

	private void SmoothTuring() {
		for (int i = 0; i < authors.size(); i++) {
			ArrayList<String> bigramstring = bigramstringarr[i];
			ArrayList<Integer> bigramvalue = bigramvaluearr[i];
			int[] Nm = new int[6];
			// Smooth the bigrams.
			int N = unigramstringarr[i].size();
			Nm[0] = N ^ 2 - bigramstring.size();
			for (int j = 0; j < bigramvalue.size(); j++) {
				if (bigramvalue.get(j) < 6) {
					Nm[bigramvalue.get(j)]++;
				}
			}
			// Fill out smoothedfreq
			for (int j = 0; j < 5; j++) {
				if (Nm[j + 1] != 0 && Nm[j] != 0) {
					smoothedfreq[i][j] = (j + 1) * Nm[j + 1] / Nm[j];
				} else {
					smoothedfreq[i][j] = j;
				}
			}
		}
	}

	private ArrayList<String> guessAuthors(String file) {
		ArrayList<String> guesses = new ArrayList<String>();
		try {
			FileReader fread = new FileReader(file);
			BufferedReader br = new BufferedReader(fread);
			String strLine;
			while ((strLine = br.readLine()) != null) {
				String[] array = strLine.split("\\s+");
				ArrayList<String> targetlist = new ArrayList<String>();
				for (int i = 0; i < array.length; i++) {
					targetlist.addAll(breakUpPunctuation(array[i]));
				}
				guesses.add(bestAuthor(targetlist));
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		return guesses;
	}

	private String bestAuthor(ArrayList<String> input) {
		String bestAuthor = "";
		Double bestPerplexity = 0.0;
		for (int i = 0; i < authors.size(); i++) {
			int N = input.size() - 1;
			Double logP = calcProbability(input, i);
			Double Perplexity = Math.pow(10, ((-1) / N * logP));
			if (bestPerplexity == 0.0 || (bestPerplexity != 0.0 && Perplexity < bestPerplexity)) {
				bestAuthor = authors.get(i);
				bestPerplexity = Perplexity;
			}
		}
		return bestAuthor;
	}

	private Double calcProbability(ArrayList<String> input, int i) {
		int index = unigramstringarr[i].indexOf(input.get(1));
		Double logP = Math.log10(unigramvaluearr[i].get(index) / totalunifreq.get(i));
		for (int j = 2; j < input.size() - 1; j++) {
			index = bigramstringarr[i].indexOf(input.get(j) + " " + input.get(j + 1));
			int unsmooth = bigramvaluearr[i].get(index);
			Double smooth;
			if (unsmooth < 5) {
				smooth = (double) smoothedfreq[i][unsmooth];
			} else {
				smooth = (double) unsmooth;
			}
			logP = logP + Math.log10(smooth / unigramvaluearr[i].get(unigramstringarr[i].indexOf(input.get(j))));
		}
		return logP;
	}

	private static ArrayList<String> breakUpPunctuation(String target) {
		ArrayList<String> output = new ArrayList<String>();
		String prev = "";
		// We want to avoid treating HTML tags as words
		if (target.length() > 0 && target.charAt(0) == '<') {
			return output;
		}
		for (int ptr = 0; ptr < target.length(); ptr++) {
			if (target.charAt(ptr) == '.' || target.charAt(ptr) == '!' || target.charAt(ptr) == '?'
					|| target.charAt(ptr) == ':' || target.charAt(ptr) == ';') {
				if (prev.length() > 0) {
					output.add(prev);
					prev = "";
				}
				output.add(target.substring(ptr, ptr + 1));
			} else if (target.charAt(ptr) != '(' && target.charAt(ptr) != ')' && target.charAt(ptr) != '\''
					&& target.charAt(ptr) != '\"' && target.charAt(ptr) != ',') {
				prev = prev + target.substring(ptr, ptr + 1);
			}
		}
		if (prev.length() > 0) {
			output.add(prev);
			prev = "";
		}
		return output;
	}
}
